#!/usr/bin/env python
# -*- coding:utf-8 -*-

import os
import sys
import time

from bsddb import db

def DEBUG(format, *args) : print format % args

class   KeyExistError(Exception) : pass

class   Yoghurt(object) :
    def __init__(self, dbpath, monolithic=True) :
        self.db_path = dbpath
        self.db_env = None

        self.db_type = db.DB_BTREE
        self.db_mode = 0660
        self.db_name = None
        self.db_open_flags = 0
        self.db_set_flags = 0

        self.monolithic = monolithic

        self.dbh = None
        self.__inited = False
        # __init__()

    def __del__(self) :
        self.__exit()
        # __del__()

    def __init(self) :
        if self.__inited : return True

        self.dbh = db.DB(self.db_env)
        self.dbh.set_flags(self.db_set_flags)
        #self.dbh.set_pagesize(8192)     # default 4k, 10000records: 72s --> 27s
        #self.dbh.set_bt_minkey(3)       # bad --> good :  4 << 3 < 2
        self.dbh.set_pagesize(16384)     # default 4k, 10000records: 72s --> 18s
        self.dbh.set_bt_minkey(8)        # bad --> good :  8 << 6 < 7 < 4 < 2
        if self.db_name:
            self.dbh.open(self.db_path, self.db_name, self.db_type,
                            self.db_open_flags|db.DB_CREATE, self.db_mode)
        else:
            self.dbh.open(self.db_path,
                            mode = self.db_mode,
                            dbtype = self.db_type,
                            flags = self.db_open_flags|db.DB_CREATE)

        self.__inited = True
        # __init()

    def __exit(self) :
        if not self.__inited : return

        self.dbh.sync()
        self.dbh.close()
        self.dbh = None
        # __exit()

    def insert(self, key, value)    :
        '''insert (key, value) pair

        if key-exists, just return
        else insert them
        '''
        if not self.__inited : self.__init()

        try:
            self.dbh.put(key, value, flags=db.DB_NOOVERWRITE)
        except db.DBKeyExistError, ex:
            #assert ex[0] == db.DB_KEYEXIST
            DEBUG('KeyExistError: %s', ex)
            raise KeyExistError, ex[1]
        except :
            DEBUG('Exception: %s', sys.exc_info())
            raise
        # insert()

    def erase(self, key)    :
        '''erase record from DB by key

        if key exists, erase it
        else just return
        '''
        if not self.__inited : self.__init()

        try :
            self.dbh.delete(key)
        except db.DBNotFoundError, ex:
            #assert ex[0] == db.DB_NOTFOUND
            DEBUG('NotFoundError: %s', ex)
            pass
        except :
            DEBUG('Exception: %s', sys.exc_info())
            raise
        # erase()

    def find(self, key) :
        '''find record from DB by key

        if key exists, return value
        else return None
        '''
        if not self.__inited : self.__init()

        try :
            value = self.dbh.get(key)
            return value
        except db.DBNotFoundError, ex:
            #assert ex[0] == db.DB_NOTFOUND
            DEBUG('NotFoundError: %s', ex)
            return None
        except :
            DEBUG('Exception: %s', sys.exc_info())
            raise
        # find()

    # class Yoghurt

#######################################################################

DASH = '-'

def makeData(key) :
    return DASH.join([key] * 32)
    # makeData()

def testKeyExist(yog) :
    key = 'abcd'
    value = makeData(key)
    yog.insert(key, value)
    yog.insert(key, value)
    # testKeyExist()

def testInsertKeys(yog, num_records) :
    for x in range(num_records/2):
        key = '%032d' % (num_records - x)  # insert keys in reverse order
        data = makeData(key)
        yog.insert(key, data)

    for x in range(num_records/2-1):
        key = '%032d' % x  # and now some in forward order
        data = makeData(key)
        yog.insert(key, data)

    # testInsertKeys()

def testFindByKeys(yog, num_records) :
    for x in range(num_records/2):
        key = '%032d' % (num_records - x)  # find keys in reverse order
        data = makeData(key)
        value = yog.find(key)
        assert data == value

    for x in range(num_records/2-1):
        key = '%032d' % x  # and now some in forward order
        data = makeData(key)
        value = yog.find(key)
        assert data == value

    # testFindByKeys()

def main() :
    db_path = '/tmp/yoghurt.db'
    os.remove(db_path)

    yog = Yoghurt(db_path)
    del yog

    #testKeyExist(yog)

    num_records = 100000

    time_start = time.time()
    yog = Yoghurt(db_path)
    testInsertKeys(yog, num_records)
    del yog
    time_stop = time.time()
    time_last = time_stop - time_start
    tps = num_records / time_last
    DEBUG('INSERT: records:%d time-last:%.2f seconds tps:%.2f', num_records, time_last, tps)

    time_start = time.time()
    yog = Yoghurt(db_path)
    testFindByKeys(yog, num_records)
    del yog
    time_stop = time.time()
    time_last = time_stop - time_start
    tps = num_records / time_last
    DEBUG('  FIND: records:%d time-last:%.2f seconds tps:%.2f', num_records, time_last, tps)

    # main()

if __name__ == '__main__' :
    main()

